#include <iostream>
#include <chrono>
#include <vector>
#include <numeric>
#include <execution>
#include <coroutine>
#include <thread>

#include "MyCoroutine.h"

// 计算平方和的任务函数
int64_t square_sum(int64_t start, int64_t end)
{
    int64_t sum = 0;
    for (int i = start; i <= end; ++i) {
        sum += i * i;
    }

    return sum;
}

// 多线程版本
void multi_threaded_square_sum(int64_t nums, int num_threads) {
    std::vector<std::thread> threads;
    std::vector<int64_t> results(num_threads, 0);
    int64_t chunk_size = nums / num_threads;

    auto start_time = std::chrono::high_resolution_clock::now();

    for (int i = 0; i < num_threads; ++i)
    {
        int64_t start = i * chunk_size;
        int64_t end = (i == num_threads - 1) ? nums : start + chunk_size;
        threads.emplace_back([&, start, end, i] {
            results[i] = square_sum(start, end);
        });
    }

    for (auto& t : threads)
    {
        if (t.joinable())
        {
            t.join();
        }
    }

    int64_t total = std::accumulate(results.begin(), results.end(), 0);

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> diff = end_time - start_time;
    std::cout << "Multi-threaded result: " << total << ", Time taken: " << diff.count() << " s\n";
}

// 协程版本的任务函数
MyCoroutine coroutine_square_sum(int64_t start, int64_t end)
{
    co_return square_sum(start, end);
}

// 多协程版本
void multi_coroutine_square_sum(int64_t nums, int num_tasks) {
    std::vector<MyCoroutine> simpleCoroutines;
    std::vector<int64_t> results(num_tasks, 0);
    int64_t chunk_size = nums / num_tasks;

    auto start_time = std::chrono::high_resolution_clock::now();

    for (int i = 0; i < num_tasks; ++i)
    {
        int64_t start = i * chunk_size;
        int64_t end = (i == num_tasks - 1) ? nums : start + chunk_size;
        simpleCoroutines.emplace_back(coroutine_square_sum(start, end));
    }

    for (int i = 0; i < num_tasks; ++i)
    {
        simpleCoroutines[i].Resume();
    }

    for (int i = 0; i < num_tasks; ++i)
    {
        while (!simpleCoroutines[i].Done())
        {
            simpleCoroutines[i].Resume();
        }

        results.push_back(simpleCoroutines[i].GetValue());
    }

    int64_t total = std::accumulate(results.begin(), results.end(), 0LL);

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> diff = end_time - start_time;
    std::cout << "Multi-coroutine result: " << total << ", Time taken: " << diff.count() << " s\n";
}

void multi_coroutine_square_sum(int64_t nums, int num_threads, int num_coroutines_per_thread) {
    std::vector<std::thread> threads;
    std::vector<std::vector<MyCoroutine>> simpleCoroutines(num_threads);
    std::vector<std::vector<int64_t>> results(num_threads);
    int64_t chunk_size = nums / num_threads;

    auto start_time = std::chrono::high_resolution_clock::now();

    for (int i = 0; i < num_threads; ++i) {
        int64_t start = i * chunk_size;
        for (int j = 0; j < num_coroutines_per_thread; ++j) {
            simpleCoroutines[i].emplace_back(coroutine_square_sum(start + j * (chunk_size / num_coroutines_per_thread),
                                             start + (j + 1) * (chunk_size / num_coroutines_per_thread)));
            results[i].push_back(0);
        }
    }

    for (int i = 0; i < num_threads; ++i) {
        threads.emplace_back([&, i]() {
            for (int j = 0; j < num_coroutines_per_thread; ++j) {
                simpleCoroutines[i][j].Resume();
                while (!simpleCoroutines[i][j].Done()) {
                    simpleCoroutines[i][j].Resume();
                }
                results[i][j] = simpleCoroutines[i][j].GetValue();
            }
        });
    }

    for (auto& t : threads) {
        t.join();
    }

    int64_t total = 0;
    for (const auto& res : results) {
        total += std::accumulate(res.begin(), res.end(), 0LL);
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> diff = end_time - start_time;
    std::cout << "Multi-thread, multi-coroutine result: " << total << ", Time taken: " << diff.count() << " s\n";
}

void test()
{
    int64_t nums = 1000000000;
    int threadNums = 8;

    multi_threaded_square_sum(nums, threadNums);
    multi_coroutine_square_sum(nums, threadNums);
    multi_coroutine_square_sum(nums, 4, 2);
}

int main()
{
    test();

    return 0;
}